import chisel3._
import chisel3.util._

class MyDCache(val SIZE: Int, val LINE_SIZE: Int, val WAY_NUM: Int, val UseLRU: Boolean)
    extends Module {
  // ============ 常用Width定义 ===========
  val INDEX_WIDTH   = log2Ceil(SIZE)
  val OFFSET_WIDTH  = log2Ceil(LINE_SIZE)
  val WAY_NUM_WIDTH = log2Ceil(WAY_NUM)
  val LINE_BIT_SIZE = LINE_SIZE * 8
  val LINE_WORD_NUM = LINE_SIZE / 4

  println("Data Cache Size: " + SIZE)
  println("Data Cache Line Size   : " + LINE_WORD_NUM + " Word | " + LINE_SIZE + " Byte | " + LINE_BIT_SIZE + " Bit")
  println("Data Cache Way Number  : " + WAY_NUM)

  assert(
    (INDEX_WIDTH + OFFSET_WIDTH <= 12).asBool,
    "Index width + offset width should be <= 12, but given is %d\n",
    (INDEX_WIDTH + OFFSET_WIDTH).asUInt
  )

  // ============ 外部接口(io)定义 ===========
  val io = IO(new Bundle {
    val valid     = Input(Vec(2, Bool()))
    val op        = Input(Vec(2, Bool()))
    val tag       = Input(Vec(2, UInt(20.W)))
    val index     = Input(Vec(2, UInt(INDEX_WIDTH.W)))
    val offset    = Input(Vec(2, UInt(OFFSET_WIDTH.W)))
    val size      = Input(Vec(2, UInt(2.W)))
    val wstrb     = Input(Vec(2, UInt(4.W)))
    val wdata     = Input(Vec(2, UInt(32.W)))
    val addr_ok   = Output(Vec(2, Bool()))
    val data_ok   = Output(Vec(2, Bool()));
    val rdata     = Output(Vec(2, UInt(32.W)))
    val uncached  = Input(Vec(2, Bool()))

    val cacop_en         = Input(Vec(2, Bool()))
    val cacop_op         = Input(Vec(2, UInt(2.W)))
    val tlb_excp_cancel_req = Input(Bool())

    // AXI接口
    val rd_req      = Output(Bool())
    val rd_type     = Output(UInt(3.W))
    val rd_addr     = Output(UInt(32.W))
    val rd_rdy      = Input(Bool())
    val ret_valid   = Input(Bool())
    val ret_last    = Input(Bool())
    val ret_data    = Input(UInt(32.W))
    val wr_req      = Output(Bool())
    val wr_type     = Output(UInt(3.W))
    val wr_addr     = Output(UInt(32.W))
    val wr_wstrb    = Output(UInt(4.W))
    val wr_data     = Output(UInt(LINE_BIT_SIZE.W))
    val wr_rdy      = Input(Bool())
  })

  // ========== 计数器统计 ===========
  val total_req_counter   = RegInit(0.U(32.W))
  val total_miss_counter  = RegInit(0.U(32.W))
  val dcache_counter      = RegInit(0.U(32.W))
  dcache_counter := dcache_counter + 1.U

  // ========== 存储结构 ===========
  val Dirty = RegInit(VecInit.fill(SIZE, WAY_NUM)(false.B))
  val DataBank = Array.fill(WAY_NUM)(Module(new DataRAM(SIZE, LINE_BIT_SIZE)).io)
  val TagvBank = Array.fill(WAY_NUM)(Module(new TagvRAM(SIZE, 20)).io)
  val Valid = RegInit(VecInit(Seq.fill(SIZE)(0.U(WAY_NUM.W))))

  // ========== LRU替换策略实现 ===========
  val LRUInfo = if (UseLRU) {
    Some(RegInit(VecInit.fill(SIZE)(VecInit.tabulate(WAY_NUM)(i => i.U(WAY_NUM_WIDTH.W)))))
  } else None

  // ========== 命中判断 ===========
  val way_hit         = Wire(Vec(2, Vec(WAY_NUM, Bool())))
  val cache_hit       = Wire(Vec(2, Bool()))
  val cache_hit_way   = Wire(Vec(2, UInt(WAY_NUM_WIDTH.W)))

  // ========== 替换路相关 ==========
  val replaced_way_buffer = RegInit(VecInit(Seq.fill(2)(0.U(WAY_NUM_WIDTH.W))))
  val replaced_index_buffer = RegInit(VecInit(Seq.fill(2)(0.U(INDEX_WIDTH.W))))
  val replaced_way      = Wire(Vec(2, UInt(WAY_NUM_WIDTH.W)))
  val replaced_data     = Wire(Vec(2, UInt(LINE_BIT_SIZE.W)))
  val replaced_tag      = Wire(Vec(2, UInt(20.W)))
  val replaced_index    = Wire(Vec(2, UInt(INDEX_WIDTH.W)))
  val replaced_addr     = Wire(Vec(2, UInt(32.W)))
  val replaced_dirty    = Wire(Vec(2, Bool()))
  val replaced_valid    = Wire(Vec(2, Bool()))

  val need_write_axi = RegInit(VecInit(Seq.fill(2)(false.B)))
  val need_read_axi  = RegInit(VecInit(Seq.fill(2)(false.B)))

  // ========== 读/写数据通道 ===========
  val read_data      = Wire(Vec(WAY_NUM, UInt(LINE_BIT_SIZE.W)))
  val target_data0   = Wire(Vec(LINE_WORD_NUM, UInt(32.W)))
  val target_data1   = Wire(Vec(LINE_WORD_NUM, UInt(32.W)))
  val read_tagv      = Wire(Vec(WAY_NUM, UInt(21.W)))
  val read_tag       = Wire(Vec(WAY_NUM, UInt(20.W)))
  val read_valid     = Wire(Vec(WAY_NUM, Bool()))

  // ========== 请求状态/缓冲 ==========
  val receive_request = Wire(Vec(2, Bool()))
  val uncached_write  = Wire(Vec(2, Bool()));
  val uncached_read   = Wire(Vec(2, Bool()));
  val uncached_en     = Wire(Vec(2, Bool()));
  val read_hit        = Wire(Vec(2, Bool()));
  val write_hit       = Wire(Vec(2, Bool()));
  val cache_unbussy   = Wire(Vec(2, Bool()));

  // ========== AXI仲裁 ==========
  val rd_arb = Module(new Arbiter(Bool(), 2))
  val rd_req_vec = Wire(Vec(2, Bool()))
  val wr_req_vec = Wire(Vec(2, Bool()))

  // ========== Cache状态机 ============
  object CacheState extends ChiselEnum {
    val IDLE, LOOKUP, WRITEHIT, MISS, REPLACE, REFILL, RESPOND = Value
  }
  val main_state = RegInit(VecInit(Seq.fill(2)(CacheState.IDLE)))

  // 状态机便捷信号
  val is_idle     = main_state.map(_ === CacheState.IDLE)
  val is_lookup   = main_state.map(_ === CacheState.LOOKUP)
  val is_writehit = main_state.map(_ === CacheState.WRITEHIT)
  val is_miss     = main_state.map(_ === CacheState.MISS)
  val is_replace  = main_state.map(_ === CacheState.REPLACE)
  val is_refill   = main_state.map(_ === CacheState.REFILL)
  val is_respond  = main_state.map(_ === CacheState.RESPOND)

  // ========== 请求缓冲器 ==========
  val request_cacop_en_buffer = RegInit(VecInit(Seq.fill(2)(false.B)))
  val request_cacop_op_buffer = RegInit(VecInit(Seq.fill(2)(0.U(2.W))))
  val request_op_buffer       = RegInit(VecInit(Seq.fill(2)(false.B)))
  val request_tag_buffer      = RegInit(VecInit(Seq.fill(2)(0.U(20.W))))
  val request_index_buffer    = RegInit(VecInit(Seq.fill(2)(0.U(INDEX_WIDTH.W))))
  val request_offset_buffer   = RegInit(VecInit(Seq.fill(2)(0.U(OFFSET_WIDTH.W))))
  val request_uncached_buffer = RegInit(VecInit(Seq.fill(2)(false.B)))
  val target_word_id          = RegInit(VecInit(Seq.fill(2)(0.U((OFFSET_WIDTH-2).W))))
  val request_wdata_buffer    = RegInit(VecInit(Seq.fill(2)(0.U(32.W))))
  val request_wstrb_buffer    = RegInit(VecInit(Seq.fill(2)(0.U(4.W))))
  val request_size_buffer     = RegInit(VecInit(Seq.fill(2)(0.U(2.W))))
  
  val request_addr = Wire(Vec(2, UInt(32.W)))
  val request_line_addr = Wire(Vec(2, UInt(32.W)))
  val normal_request = Wire(Vec(2, Bool()))

  for (i <- 0 until 2) {
    request_addr(i) := Cat(Mux(is_lookup(i), io.tag(i), request_tag_buffer(i)), 
                          request_index_buffer(i), request_offset_buffer(i))
    request_line_addr(i) := Cat(Mux(is_lookup(i), io.tag(i), request_tag_buffer(i)), 
                              request_index_buffer(i), 0.U(OFFSET_WIDTH.W))
  }

  val ret_data_buffer = Reg(Vec(2, Vec(LINE_WORD_NUM, UInt(32.W))))
  val miss_data_ret_cnt = RegInit(VecInit(Seq.fill(2)(0.U((OFFSET_WIDTH - 2).W))))

  // ========== cacop指令处理 ==========
  val cacop_op_0 = Wire(Vec(2, Bool()))
  val cacop_op_1 = Wire(Vec(2, Bool()))
  val cacop_op_2 = Wire(Vec(2, Bool()));
  val cacop_way  = Wire(Vec(2, UInt(WAY_NUM_WIDTH.W)))

  for (i <- 0 until 2) {
    cacop_op_0(i) := request_cacop_en_buffer(i) && request_cacop_op_buffer(i) === 0.U
    cacop_op_1(i) := request_cacop_en_buffer(i) && request_cacop_op_buffer(i) === 1.U
    cacop_op_2(i) := request_cacop_en_buffer(i) && request_cacop_op_buffer(i) === 2.U
    cacop_way(i)  := request_addr(i)(WAY_NUM_WIDTH-1, 0)
  }

  // ========== 替换路选择 ==========
  for (i <- 0 until 2) {
    replaced_way(i) := 
 
        Mux(cacop_op_2(i) && cache_hit(i), cache_hit_way(i), 
          Mux(cacop_op_0(i) || cacop_op_1(i), cacop_way(i), 
            dcache_counter(WAY_NUM_WIDTH-1, 0) + i.U))
      
    
    replaced_data(i)   := read_data(replaced_way(i))
    replaced_dirty(i)  := Dirty(replaced_index(i))(replaced_way(i))
    replaced_index(i)  := request_index_buffer(i)
    replaced_tag(i)    := read_tag(replaced_way(i))
    replaced_valid(i)  := read_valid(replaced_way(i))
    replaced_addr(i)   := Cat(replaced_tag(i), replaced_index(i), 0.U(OFFSET_WIDTH.W))
  }

  // ========== Tag和Valid读取 ==========
  for (i <- 0 until WAY_NUM) {
    read_tag(i)   := read_tagv(i)(20, 1)
    read_valid(i) := read_tagv(i)(0)
  }

  // ========== 命中判定 ==========
  for (i <- 0 until WAY_NUM) {
    way_hit(0)(i) := (read_tag(i) === io.tag(0)) && read_valid(i)
    way_hit(1)(i) := (read_tag(i) === io.tag(1)) && read_valid(i)
  }

  for (i <- 0 until 2) {
    cache_hit(i)       := way_hit(i).reduce(_||_) && !io.tlb_excp_cancel_req
    cache_hit_way(i)   := OHToUInt(way_hit(i))
    uncached_en(i)     := Mux(is_lookup(i), io.uncached(i), request_uncached_buffer(i)) && !request_cacop_en_buffer(i)
    uncached_read(i)   := uncached_en(i) && !request_op_buffer(i)
    uncached_write(i)  := uncached_en(i) && request_op_buffer(i)
    normal_request(i)  := !Mux(is_lookup(i), io.uncached(i), request_uncached_buffer(i)) && !request_cacop_en_buffer(i)
    read_hit(i)        := cache_hit(i) && !request_op_buffer(i) && normal_request(i)
    write_hit(i)       := cache_hit(i) && request_op_buffer(i) && normal_request(i)
    cache_unbussy(i)   := is_idle(i) || (is_lookup(i) && cache_hit(i) && normal_request(i)) || is_respond(i) || 
                         (is_replace(i) && !need_read_axi(i))
    receive_request(i) := cache_unbussy(i) && (io.valid(i) || io.cacop_en(i))
    io.addr_ok(i)      := cache_unbussy(i)
    io.data_ok(i)      := is_lookup(i) && (cache_hit(i) || request_op_buffer(i)) && !uncached_en(i) || is_respond(i)
  }

  // ========== 数据读取和组织 ==========
  target_data0 := read_data(cache_hit_way(0)).asTypeOf(target_data0)
  target_data1 := read_data(cache_hit_way(1)).asTypeOf(target_data1)
  
  for (i <- 0 until 2) {
    io.rdata(i) := Mux(is_respond(i), 
                      ret_data_buffer(i)(target_word_id(i)), 
                      Mux(i.U === 0.U, target_data0(target_word_id(i)), target_data1(target_word_id(i))))
  }

  // ========== AXI接口控制 ==========
  for (i <- 0 until 2) {
    rd_req_vec(i) := is_replace(i) && need_read_axi(i)
    wr_req_vec(i) := is_miss(i) && need_write_axi(i)
  }

  rd_arb.io.in(0).valid := rd_req_vec(0)
  rd_arb.io.in(1).valid := rd_req_vec(1)
  rd_arb.io.in(0).bits := true.B
  rd_arb.io.in(1).bits := true.B

  val rd_grant = rd_arb.io.out.valid && rd_arb.io.out.bits
  val rd_port_sel = rd_arb.io.chosen

  io.rd_req := rd_arb.io.out.valid
  io.rd_type := Mux(rd_port_sel, rd_type_buffer(1), rd_type_buffer(0))
  io.rd_addr := Mux(rd_port_sel, rd_addr_buffer(1), rd_addr_buffer(0))

  val wr_arb = Module(new Arbiter(Bool(), 2))
  wr_arb.io.in(0).valid := wr_req_vec(0)
  wr_arb.io.in(1).valid := wr_req_vec(1)
  wr_arb.io.in(0).bits := true.B
  wr_arb.io.in(1).bits := true.B

  val wr_grant = wr_arb.io.out.valid && wr_arb.io.out.bits
  val wr_port_sel = wr_arb.io.chosen

  io.wr_req := wr_arb.io.out.valid
  io.wr_type := Mux(wr_port_sel, wr_type_buffer(1), wr_type_buffer(0))
  io.wr_addr := Mux(wr_port_sel, wr_addr_buffer(1), wr_addr_buffer(0))
  io.wr_wstrb := Mux(wr_port_sel, wr_wstrb_buffer(1), wr_wstrb_buffer(0))
  io.wr_data := Mux(wr_port_sel, wr_data_buffer(1), wr_data_buffer(0))

  // ========== 数据合并 ==========
  for (i <- 0 until 2) {
    val ret_data_v = io.ret_data.asTypeOf(Vec(4, UInt(8.W)))
    val target_data_v = Mux(i.U === 0.U, 
                          target_data0(target_word_id(0)).asTypeOf(Vec(4, UInt(8.W))),
                          target_data1(target_word_id(1)).asTypeOf(Vec(4, UInt(8.W))))
    val request_wdata_v = request_wdata_buffer(i).asTypeOf(Vec(4, UInt(8.W)))
    
    val merged_data = Wire(Vec(4, UInt(8.W)))
    for (j <- 0 until 4) {
      merged_data(j) := Mux(request_wstrb_buffer(i)(j), request_wdata_v(j), ret_data_v(j))
    }
    
    when(is_refill(i) && io.ret_valid) {
      ret_data_buffer(i)(miss_data_ret_cnt(i)) := merged_data.asUInt
    }
  }

  // ========== 数据银行控制 ==========
  val conflict_data = Reg(Vec(WAY_NUM, UInt(LINE_BIT_SIZE.W)))
  
  for (p <- 0 until 2) {
    for (i <- 0 until WAY_NUM) {
      DataBank(i).clka := clock
      DataBank(i).clkb := clock
      
      // 写端口A
      DataBank(i).addra := Mux(write_hit(p) && way_hit(p)(i), 
                              request_index_buffer(p),
                              replaced_index_buffer(p))
      DataBank(i).dina := Mux(is_lookup(p), 
                            Mux(p.U === 0.U, target_data0.asUInt, target_data1.asUInt), 
                            ret_data_buffer(p).asUInt)
      DataBank(i).wea := is_lookup(p) && write_hit(p) && way_hit(p)(i) || 
                        is_respond(p) && normal_request(p) && (i.U === replaced_way_buffer(p))
      
      // 读端口B
      DataBank(i).addrb := io.index(p)
      read_data(i) := DataBank(i).doutb
      
      // 冲突检测
      when((request_index_buffer(0) === io.index(p) || request_index_buffer(1) === io.index(p)) && DataBank(i).wea) {
        conflict_data(i) := DataBank(i).dina
      }.otherwise {
        conflict_data(i) := 0.U
      }
    }
  }

  // ========== Tag和Valid写入 ==========
  val valid_in = Wire(Vec(WAY_NUM, UInt(1.W)))
  
  for (p <- 0 until 2) {
    for (i <- 0 until WAY_NUM) {
      read_tagv(i) := Cat(TagvBank(i).douta, Valid(request_index_buffer(p)))
      TagvBank(i).clka := clock
      TagvBank(i).addra := Mux(receive_request(p), io.index(p), request_index_buffer(p))
      TagvBank(i).dina := Mux(request_cacop_en_buffer(p), 0.U(20.W), request_tag_buffer(p))
      TagvBank(i).wea := is_miss(p) && !uncached_en(p) && ((i.U === replaced_way_buffer(p)) || request_cacop_en_buffer(p))
      valid_in(p)(i) := Mux(TagvBank(i).wea, !request_cacop_en_buffer(p), Valid(request_index_buffer(p))(i))
    }
    Valid(request_index_buffer(p)) := valid_in(p).asUInt
  }

  // ========== 状态机主逻辑 ==========
  for (i <- 0 until 2) {
    switch(main_state(i)) {
      is(CacheState.IDLE) {
        when(receive_request(i)) {
          request_cacop_en_buffer(i) := io.cacop_en(i)
          request_cacop_op_buffer(i) := io.cacop_op(i)
          request_op_buffer(i) := io.op(i)
          request_tag_buffer(i) := io.tag(i)
          request_index_buffer(i) := io.index(i)
          request_offset_buffer(i) := io.offset(i)
          request_size_buffer(i) := io.size(i)
          request_wstrb_buffer(i) := io.wstrb(i)
          request_wdata_buffer(i) := io.wdata(i)
          target_word_id(i) := io.offset(i)(OFFSET_WIDTH-1, 2)
          main_state(i) := CacheState.LOOKUP
        }
      }
      is(CacheState.LOOKUP) {
        total_req_counter := total_req_counter + 1.U
        request_uncached_buffer(i) := io.uncached(i)

        when(io.tlb_excp_cancel_req) {
          main_state(i) := CacheState.IDLE
        }.elsewhen(!cache_hit(i) || request_cacop_en_buffer(i) || uncached_en(i)) {
          when(uncached_en(i)) { target_word_id(i) := 0.U }
          replaced_way_buffer(i) := replaced_way(i)
          replaced_index_buffer(i) := replaced_index(i)
          need_write_axi(i) := (cacop_op_2(i) && cache_hit(i) || cacop_op_1(i) || cacop_op_0(i) || normal_request(i)) && 
                              replaced_dirty(i) && replaced_valid(i) || uncached_write(i)
          
          // 缓冲AXI请求信息
          rd_type_buffer(i) := Mux(uncached_en(i), Cat(0.U(1.W), request_size_buffer(i)), "b100".U)
          rd_addr_buffer(i) := Mux(uncached_en(i), request_addr(i), request_line_addr(i))
          wr_type_buffer(i) := Mux(uncached_en(i), Cat(0.U(1.W), request_size_buffer(i)), "b100".U)
          wr_addr_buffer(i) := Mux(uncached_en(i), request_addr(i), replaced_addr(i))
          wr_data_buffer(i) := Mux(uncached_en(i), request_wdata_buffer(i), replaced_data(i))
          wr_wstrb_buffer(i) := Mux(uncached_en(i), request_wstrb_buffer(i), "b1111".U)
          
          main_state(i) := Mux(cacop_op_2(i) && !cache_hit(i), CacheState.IDLE, CacheState.MISS)
          total_miss_counter := total_miss_counter + 1.U
        }.elsewhen(receive_request(i)) {
          // 连续命中时提前接收新请求
          request_cacop_en_buffer(i) := io.cacop_en(i)
          request_cacop_op_buffer(i) := io.cacop_op(i)
          request_op_buffer(i) := io.op(i)
          request_tag_buffer(i) := io.tag(i)
          request_index_buffer(i) := io.index(i)
          request_offset_buffer(i) := io.offset(i)
          request_size_buffer(i) := io.size(i)
          request_wstrb_buffer(i) := io.wstrb(i)
          request_wdata_buffer(i) := io.wdata(i)
          target_word_id(i) := io.offset(i)(OFFSET_WIDTH-1, 2)
          main_state(i) := CacheState.LOOKUP
        }.otherwise {
          main_state(i) := CacheState.IDLE
        }
        
        when(write_hit(i)) {
          Dirty(request_index_buffer(i))(cache_hit_way(i)) := true.B
          if (UseLRU) {
            // 更新LRU信息
            for (j <- 0 until WAY_NUM) {
              when(j.U =/= cache_hit_way(i)) {
                LRUInfo.get(request_index_buffer(i))(j) := Mux(
                  LRUInfo.get(request_index_buffer(i))(j) > LRUInfo.get(request_index_buffer(i))(cache_hit_way(i)),
                  LRUInfo.get(request_index_buffer(i))(j) - 1.U,
                  LRUInfo.get(request_index_buffer(i))(j)
                )
              }
            }
            LRUInfo.get(request_index_buffer(i))(cache_hit_way(i)) := (WAY_NUM - 1).U
          }
        }
      }
      is(CacheState.MISS) {
        when(io.wr_rdy || !need_write_axi(i)) {
          need_read_axi(i) := uncached_read(i) || normal_request(i)
          main_state(i) := CacheState.REPLACE
        }
      }
      is(CacheState.REPLACE) {
        when(need_read_axi(i)) {
          when(io.rd_rdy) {
            main_state(i) := CacheState.REFILL
            miss_data_ret_cnt(i) := 0.U
            ret_data_buffer(i) := (0.U).asTypeOf(ret_data_buffer(i))
          }
        }.otherwise {
          when(receive_request(i)) {
            request_cacop_en_buffer(i) := io.cacop_en(i)
            request_cacop_op_buffer(i) := io.cacop_op(i)
            request_op_buffer(i) := io.op(i)
            request_tag_buffer(i) := io.tag(i)
            request_index_buffer(i) := io.index(i)
            request_offset_buffer(i) := io.offset(i)
            request_size_buffer(i) := io.size(i)
            request_wstrb_buffer(i) := io.wstrb(i)
            request_wdata_buffer(i) := io.wdata(i)
            target_word_id(i) := io.offset(i)(OFFSET_WIDTH-1, 2)
            main_state(i) := CacheState.LOOKUP
          }.otherwise {
            main_state(i) := CacheState.IDLE
          }
        }
      }
      is(CacheState.REFILL) {
        when(io.ret_valid) {
          miss_data_ret_cnt(i) := miss_data_ret_cnt(i) + 1.U
          when(io.ret_last) {
            when(normal_request(i)) {
              Dirty(replaced_index_buffer(i))(replaced_way_buffer(i)) := request_op_buffer(i)
              Valid(replaced_index_buffer(i))(replaced_way_buffer(i)) := true.B
              TagvBank(replaced_way_buffer(i)).addra := replaced_index_buffer(i)
              TagvBank(replaced_way_buffer(i)).dina := request_tag_buffer(i)
              TagvBank(replaced_way_buffer(i)).wea := true.B
            }
            main_state(i) := CacheState.RESPOND
          }
        }
      }
      is(CacheState.RESPOND) {
        when(receive_request(i)) {
          request_cacop_en_buffer(i) := io.cacop_en(i)
          request_cacop_op_buffer(i) := io.cacop_op(i)
          request_op_buffer(i) := io.op(i)
          request_tag_buffer(i) := io.tag(i)
          request_index_buffer(i) := io.index(i)
          request_offset_buffer(i) := io.offset(i)
          request_size_buffer(i) := io.size(i)
          request_wstrb_buffer(i) := io.wstrb(i)
          request_wdata_buffer(i) := io.wdata(i)
          target_word_id(i) := io.offset(i)(OFFSET_WIDTH-1, 2)
          main_state(i) := CacheState.LOOKUP
        }.otherwise {
          main_state(i) := CacheState.IDLE
        }
      }
    }
  }
}

// 简单的仲裁器实现
class Arbiter(typ: Data, n: Int) extends Module {
  val io = IO(new Bundle {
    val in = Flipped(Vec(n, Decoupled(typ)))
    val out = Decoupled(typ)
    val chosen = Output(UInt(log2Ceil(n).W))
  })
  
  // 简单的轮转仲裁
  val grantMask = RegInit(VecInit.tabulate(n)(i => (i == 0).B))
  val grantIdx = PriorityEncoder(grantMask.zip(io.in.map(_.valid)).map { case (m, v) => m && v })
  
  io.out.valid := io.in.map(_.valid).reduce(_||_)
  io.out.bits := Mux1H(io.in.zipWithIndex.map { case (in, i) => (i.U === grantIdx) -> in.bits })
  io.chosen := grantIdx
  
  for (i <- 0 until n) {
    io.in(i).ready := io.out.ready && (grantIdx === i.U)
  }
  
  when(io.out.fire()) {
    grantMask := VecInit((0 until n).map(i => (i == (grantIdx + 1) % n).B))
  }
}